# -*- coding:utf-8 -*-
"""
@file name  : lesson-03-Linear-Regression.py
@author     : QuZhang https://gitee.com/zhang_qu/
@date       : 2020-11-21
@brief      : 一元线性回归模型
"""
import torch
import matplotlib.pylab as plt
import os
os.environ['KMP_DUPLICATE_LIB_OK']='True'
torch.manual_seed(10)

if __name__ == "__main__":
    lr = 0.05  # 学习率

    # 创建训练数据(要拟合的对象)
    # 要拟合的自变量取值x: [0,1)之间的均匀分布创建
    x = torch.rand(20, 1) * 10  # x data (tensor) , shape(20, 1)
    # 要拟合的因变量取值
    # torch.randn(20, 1)是加入高斯噪声，使得要拟合的不是标准的直线数据
    y = 2*x + (5 + torch.randn(20, 1))  # y data (tensor), shape(20, 1)

    # 构建线性回归参数
    w = torch.randn((1,), requires_grad=True)  # w 一维张量
    b = torch.zeros((1,), requires_grad=True)  # b 一维张量
    print("w:\n{}\nb:\n{}\n".format(w, b))

    # 模型学习
    # 前向传播 -> 计算损失 -> 计算梯度 -> 更新参数 ...
    for iteration in range(1000):

        # 前向传播
        wx = torch.mul(w, x)  # 广播
        y_pred = torch.add(wx, b)

        # 计算MSE loss
        loss = (0.5 * (y_pred - y) ** 2).mean()
        print("loss type：{}, loss value: {}".format(type(loss), loss))

        # 反向传播
        # 计算损失loss对w和b的梯度,会将叶子节点梯度保存好
        loss.backward()

        # 更新参数
        w.data.sub_(lr * w.grad)
        b.data.sub_(lr * b.grad)

        # 清零张量的梯度
        w.grad.zero_()
        b.grad.zero_()

        # 绘图
        if iteration % 20 == 0:
            plt.scatter(x.data.numpy(), y.data.numpy())  # 将tensor转为numpy
            plt.plot(x.data.numpy(), y_pred.data.numpy(), 'r-', lw=5)
            plt.text(2, 20, 'Loss=%.4f' % loss.data.numpy(), fontdict={'size': 20, 'color': 'red'})
            plt.xlim(1.5, 10)
            plt.ylim(8, 28)
            plt.title("Iteration: {}\nw:{} b:{}".format(iteration, w.data.numpy(), b.data.numpy()))
            plt.pause(0.5)

            if loss.data.numpy() < 1:
                break
